#include "yolo_v6_o.h"
#include "encode_dlm_v5_o.h"
#include <thread>
#include <chrono>
#include <iostream>
#include <cstdio>
#include <direct.h>
#include <vector>
#include "boost/filesystem/operations.hpp"
#include "boost/filesystem/path.hpp"
#include "boost/filesystem/fstream.hpp"

#include "common/device_utility.h"

/**
 * Enumerates all files with a specific extension in the specified path.
 *
 * @param path The directory path to enumerate.
 * @param extension The file extension to match.
 * @return A vector containing all matching file paths.
 */
std::vector<std::string> EnumFiles(
    const std::string& path,
    const std::string &extension)
{
    namespace fs = boost::filesystem;
    fs::path _path(path);
    fs::path fullpath = fs::system_complete(_path);
    std::vector<std::string> ret;

    if (!fs::exists(fullpath) && !fs::is_directory(fullpath)) {
        return ret;
    }
    fs::directory_iterator end_iter;
    for (fs::directory_iterator iter(fullpath); iter != end_iter; iter++) {
        try {
            if (fs::is_directory(*iter)) {
                continue;
            }
            else if (iter->path().extension() == extension) {
                ret.push_back(iter->path().string());
            }
        } catch (const std::exception&) {
            continue;
        }
    }
    return ret;
}

/**
 * Enumerates model configuration files and performs model conversion.
 *
 * This function iterates through all model configuration files (.yaml) in the specified directory.
 * For each configuration file, it attempts to find the corresponding model file (.onnx or .ennx),
 * then uses the YOLO v6 class to perform model conversion.
 * After successful conversion, the converted model file is renamed and moved to the destination folder.
 *
 * @param mon_folder The path to the model folder.
 * @param dest_folder The path to the destination folder.
 * @param mark_info Mark information used to generate the target file name.
 */
void enum_models(std::string& mon_folder, std::string& dest_folder, std::string & mark_info) {
    namespace fs = boost::filesystem;
    std::string gpu_model = GetGPUModel();
    std::cout <<"GPU model: " << gpu_model << std::endl;
    if (gpu_model.empty()) {
        std::cout << "Unknown GPU Model! \n";
        return;
    }

    fs::path _dest_folder(dest_folder);
    if (!fs::exists(_dest_folder)) {
        return;
    }
    std::cout << _dest_folder.string() << " exists" << std::endl;
    std::vector<std::string> cfg_files = EnumFiles(mon_folder, std::string(".yaml"));
    std::cout << "Number of cfg files: " << cfg_files.size() << std::endl;

    for (const auto & cfg_file : cfg_files) {
        std::cout << "This cfg: " << cfg_file.c_str() << std::endl;
        st_trans_model_input input;
        fs::path _cfg_path(cfg_file);
        input.cfg_path = _cfg_path.string();
        std::cout << "Input cfg path: " << input.cfg_path << std::endl;

        fs::path _model_path;
        {
            fs::path _path = _cfg_path.replace_extension(".ennx");
            if (fs::exists(_path)) {
                _model_path = _path;
            }
        }
        {
            fs::path _path = _cfg_path.replace_extension(".onnx");
            if (fs::exists(_path)) {
                _model_path = _path;
                std::cout << "Model ONNX: " << _model_path.string() << std::endl;
            }
        }
        if (_model_path.empty()) {
            continue;
        }

        fs::path _gdats_path(dest_folder + "\\" + _cfg_path.stem().string() + mark_info + MODEL_FILE_EXT);
        std::cout << "Gdats path: " << _gdats_path.string() << std::endl;
        if (fs::exists(_gdats_path)) {
            continue;
        }

        std::cout << "START===============================================\n";
        input.model_path = _model_path.string();
        input.gpu_model = gpu_model;
        input.trans_path = _gdats_path.string() + "-tmp";

        yolo_v6_o modelConverter;
        if (modelConverter.TransModel(input)) {
            std::cout << "DONE================================================\n";
            int renameResult = std::rename(
                input.trans_path.c_str(),
                _gdats_path.string().c_str());
            if (renameResult == 0) {
                std::cout << "Conversion successful and file has been renamed.\n";
            } else {
                std::cout << "Conversion successful, but file renaming failed.\n";
            }
        } else {
            std::cout << "FAIL!================================================\n";
        }
    }
}

int main(int argc, char* argv[]) {
    try {
        char cur_folder[260];
        _getcwd(cur_folder, sizeof(cur_folder));
        printf("Current working directory: \n %s\n", cur_folder);

        std::string cwd = cur_folder;
        std::string mon_folder = cwd;
        std::string dest_folder = cwd;
        std::string mark_info = "v1";

        if (argc >= 2) {
            mon_folder = argv[1];
        }
        if (argc >= 3) {
            dest_folder = argv[2];
        }
        if (argc >= 4) {
            mark_info = argv[3];
        }

        enum_models(mon_folder, dest_folder, mark_info);
        printf("Done!\n");
    } catch (...) {
        printf("Exception!\n");
    }

    return 0;
}